package com.shujia.flink.core;

import org.apache.flink.api.common.eventtime.*;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

public class Demo05WatermarkStrategy {
    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        env.setParallelism(2);

        DataStream<String> lineDS = env.socketTextStream("master", 8888);

        // 整理数据
        DataStream<Tuple2<String, Long>> wordsTSDS = lineDS.map(line -> {
            String[] split = line.split(",");
            String word = split[0];
            Long ts = Long.parseLong(split[1]);
            return Tuple2.of(word, ts);
        }, Types.TUPLE(Types.STRING, Types.LONG));

        DataStream<Tuple2<String, Long>> assignDS = wordsTSDS.assignTimestampsAndWatermarks(
                new WatermarkStrategy<Tuple2<String, Long>>() {
                    @Override
                    // 指定一个时间戳的提取策略
                    public TimestampAssigner<Tuple2<String, Long>> createTimestampAssigner(TimestampAssignerSupplier.Context context) {
                        return new TimestampAssigner<Tuple2<String, Long>>() {
                            @Override
                            public long extractTimestamp(Tuple2<String, Long> element, long recordTimestamp) {
                                return element.f1;
                            }
                        };
                    }

                    @Override
                    // 指定水位线的生成策略
                    public WatermarkGenerator<Tuple2<String, Long>> createWatermarkGenerator(WatermarkGeneratorSupplier.Context context) {
                        return new MyWatermarkGenerator(0);
                    }
                }
        );

        assignDS
                .map(t2 -> Tuple2.of(t2.f0, 1), Types.TUPLE(Types.STRING, Types.INT))
                .keyBy(kv -> kv.f0, Types.STRING)
                // 每5s触发一次
                .window(TumblingEventTimeWindows.of(Time.seconds(5)))
                // 窗口的触发条件：1、窗口内有数据 2、所有并行度的水位线都超过了窗口的结束时间
                .process(new ProcessWindowFunction<Tuple2<String, Integer>, String, String, TimeWindow>() {
                    @Override
                    // 每个窗口中的每个Key会执行一次
                    public void process(String key, ProcessWindowFunction<Tuple2<String, Integer>, String, String, TimeWindow>.Context context, Iterable<Tuple2<String, Integer>> elements, Collector<String> out) throws Exception {
                        long windowStart = context.window().getStart();
                        long windowEnd = context.window().getEnd();
                        long currentWatermark = context.currentWatermark();
                        System.out.println("窗口触发了一次，" + "当前的水位线为：" + currentWatermark + "，当前窗口的开始时间为：" + windowStart + "，当前窗口的结束时间为：" + windowEnd);

                        int sum = 0;
                        for (Tuple2<String, Integer> element : elements) {
                            sum += element.f1;
                        }
                        out.collect(key + "," + sum);
                    }
                })
                .print();

        env.execute();

    }
}

class MyWatermarkGenerator implements WatermarkGenerator<Tuple2<String, Long>> {

    long maxTimestamp = 0;
    long outOfOrderness = 0; // 以秒为单位

    public MyWatermarkGenerator(long time) {
        outOfOrderness = time;
    }

    @Override
    // 每条数据会执行一次
    public void onEvent(Tuple2<String, Long> event, long eventTimestamp, WatermarkOutput output) {
        // 取数据中最大的时间戳
        maxTimestamp = Math.max(eventTimestamp, maxTimestamp);
        System.out.println("当前线程id为：" + Thread.currentThread().getId() + "，当前线程中最大的时间戳为：" + maxTimestamp);
    }

    @Override
    // 200ms一个周期执行一次
    public void onPeriodicEmit(WatermarkOutput output) {
        // 周期性发送水位线到下游，相比来一条数据发送一次，性能会更好
        output.emitWatermark(new Watermark(maxTimestamp - outOfOrderness * 1000));
    }
}